from torch_geometric.loader import LinkNeighborLoader,ZipLoader
import torch
import copy
import torch_geometric.transforms as T
from data_package_003 import my_dataset
from torch_geometric.loader import DataLoader
from torch_geometric.data import Data
from torch_geometric.data import HeteroData

def descrition_data(_data_dict, _data):
    _descrition_data = {}
    all_list = []
    vdata = _data.edge_label_index.t().tolist()
    # print(vdata)
    for i2j in vdata:
        all_list.append((i2j, _data_dict.all_dict[i2j[0]], _data_dict.all_dict[i2j[1]]))
    return all_list



# <class 'torch_geometric.data.hetero_data.HeteroData'>
# <class 'torch_geometric.data.hetero_data.HeteroData'>
train_data_list = []

#
# print(next(iter(my_dataset)))
datas = []

aoh_list = []
atb_list = []
host_list = []
host_node_id_list = []
alarm_list = []
alarm_node_id_list = []
btree_list = []
btree_node_id_list= []
"""
Data(edge_index=[2, 1136], node_id=[389], x=[389, 768], node_type=[389], edge_type=[1136])
Data(edge_index=[2, 948], node_id=[389], x=[389, 768], node_type=[389], edge_type=[948])
Data(edge_index=[2, 1124], node_id=[389], x=[389, 768], node_type=[389], edge_type=[1124])
Data(edge_index=[2, 1170], node_id=[367], x=[367, 768], node_type=[367], edge_type=[1170])
"""

first = True
for batch in my_dataset:
    if first == True:
        first = False
        host_list.append(batch["host"].x)
        btree_list.append(batch["btree"].x)
        host_node_id_list.append(batch["host"].node_id)
        btree_node_id_list.append(batch["btree"].node_id)

    alarm_list.append(batch["alarm"].x)
    alarm_node_id_list.append(batch["alarm"].node_id)

    aoh_list.append(batch[("alarm", "on", "host")])
    atb_list.append(batch[("alarm", "to", "btree")])

data = HeteroData()
data["btree"].node_id = torch.IntTensor(torch.cat(btree_node_id_list))
data['btree'].x = torch.cat(btree_list)
# data["host"].node_id = torch.IntTensor(torch.cat(host_node_id_list))
# data['host'].x = torch.cat(host_list)
data["alarm"].node_id = torch.IntTensor(torch.cat(alarm_node_id_list))
data['alarm'].x = torch.cat(alarm_list,dim=0)
# data[('alarm', 'on', 'host')].edge_index = torch.cat([i["edge_index"] for i in aoh_list],dim=1)
data[('alarm', 'to', 'btree')].edge_index = torch.cat([i["edge_index"] for i in atb_list],dim=1)

# print("data1",data)
# del data['alarm'].num_nodes
# print("data2",data)

data  =  T.ToUndirected()(data)
# print("1",data)
# del data['btree', 'rev_to', 'alarm'].edge_label
# print("2",data)

# # 定义元路径
# from torch_geometric.transforms import AddMetaPaths
# metapaths = [[('alarm',  'btree'), ("btree", "alarm")]]
# # 创建一个转换对象，用于添加元路径
# transform = AddMetaPaths(metapaths=metapaths, drop_orig_edge_types=True, drop_unconnected_node_types=True)
# data = transform(data)
# print(data)

transform = T.RandomLinkSplit(
    num_val=0.1,
    num_test=0.1,
    disjoint_train_ratio=0.0,
    neg_sampling_ratio=0.0,
    add_negative_train_samples=False,
    edge_types=("alarm", "to", "btree"),
    rev_edge_types=("btree", "rev_to", "alarm")
)

# print(data[("alarm", "to", "btree")],"1")
train_data,val_data,test_data = transform(data)
# print(train_data[("alarm", "to", "btree")].edge_label)
# print(val_data[("alarm", "to", "btree")].edge_label)
# print(test_data[("alarm", "to", "btree")].edge_label)

# print(train_data)
# print(weighted_mse_loss(torch.FloatTensor([0.8,1.9]),torch.IntTensor([1,1]),weight))
# print(weighted_mse_loss(torch.FloatTensor([0.2,2.3]),torch.IntTensor([1,1]),weight))
# print(weighted_mse_loss(torch.FloatTensor([0.4,0.5]),torch.IntTensor([1,1]),weight))

# print("11111",train_data)

# import torch_sparse
# # 二部图操作
# from torch_geometric.nn.conv.gcn_conv import gcn_norm
# metapath = [("alarm", "to", "btree"), ("btree", "rev_to", "alarm")]
# train_data = T.AddMetaPaths(metapaths=[metapath])(train_data)
#
# # Apply normalization to filter the metapath:
# _, edge_weight = gcn_norm(
#     train_data['alarm', 'alarm'].edge_index,
#     num_nodes=train_data['alarm'].num_nodes,
#     add_self_loops=False,
# )
# edge_index = train_data['alarm', 'alarm'].edge_index[:, edge_weight > 0.002]
#
# train_data['alarm', 'metapath_0', 'alarm'].edge_index = edge_index
# val_data['alarm', 'metapath_0', 'alarm'].edge_index = edge_index
# test_data['alarm', 'metapath_0', 'alarm'].edge_index = edge_index
#
# print("22222",train_data)
# print("22222",val_data)
# print("22222",test_data)


edge_label_index = train_data[("alarm", "to", "btree")].edge_label_index
edge_label = train_data[("alarm", "to", "btree")].edge_label
# # print(edge_label.shape,"555555555555")
#
train_loader = LinkNeighborLoader(
    data=train_data,
    num_neighbors=[15, 15],
    neg_sampling_ratio=2.0,
    edge_label_index=(("alarm", "to", "btree"),edge_label_index),
    edge_label=edge_label,
    batch_size=64,
    shuffle=True
    # drop_last=True
)

HeteroDataInfo = data


